/****************************************************************************
 * Copyright (C) 2009-2010 SciTouch LLC
 * 
 * This file is part of Indigo toolkit.
 * 
 * This file may be distributed and/or modified under the terms of the
 * GNU General Public License version 3 as published by the Free Software
 * Foundation and appearing in the file LICENSE.GPL included in the
 * packaging of this file.
 * 
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 ***************************************************************************/

#include "base_cpp/crc32.h"
#include "base_c/bitarray.h"
#include "base_cpp/output.h"

#include "graph/graph_subtree_enumerator.h"
#include "graph/subgraph_hash.h"

#include "molecule/molecule.h"
#include "molecule/molecule_fingerprint.h"
#include "molecule/molecule_tautomer.h"

MoleculeFingerprintBuilder::MoleculeFingerprintBuilder (const Molecule &mol, 
                     const MoleculeFingerprintParameters &parameters):
_mol(mol),
_parameters(parameters),
TL_CP_GET(_total_fingerprint)
{
   _total_fingerprint.resize(_parameters.fingerprintSize());
   cb_fragment = 0;

   query = false;
   
   skip_ord = false;
   skip_sim = false;
   skip_tau = false;
   skip_ext = false;
   skip_ext_charge = false;

   skip_any_atoms = false;
   skip_any_bonds = false;
   skip_any_atoms_bonds = false;
}

void MoleculeFingerprintBuilder::process ()
{
   _total_fingerprint.zerofill();
   _makeFingerprint(_mol);
}

bool MoleculeFingerprintBuilder::_handleCycle (const Graph &graph,
        const Array<int> &vertices, const Array<int> &edges, void *context)
{
   MoleculeFingerprintBuilder *self = (MoleculeFingerprintBuilder *)context;
   self->_handleSubgraph(graph, vertices, edges);
   return true;
}

void MoleculeFingerprintBuilder::_handleTree (const Graph &graph,
        const int *v_mapping, const int *e_mapping, void *context)
{
   QS_DEF(Array<int>, vertices);
   QS_DEF(Array<int>, edges);

   Graph::filterVertices(graph, v_mapping, FILTER_NEQ, -1, vertices);
   Graph::filterEdges(graph, e_mapping, FILTER_NEQ, -1, edges);

   MoleculeFingerprintBuilder *self = (MoleculeFingerprintBuilder *)context;
   self->_handleSubgraph(graph, vertices, edges);
}

int MoleculeFingerprintBuilder::_maximalSubgraphCriteriaValue (const Graph &graph, 
      const int *v_mapping, const int *e_mapping, void *context)
{
   Filter filter_v(v_mapping, Filter::NEQ, -1);
   Filter filter_e(e_mapping, Filter::NEQ, -1);
   const Molecule &mol = (const Molecule &)graph;
   int ret = 0;
   
   if (Molecule::fragmentHasQueryAtoms(mol, &filter_v))
      ret |= 1;

   if (Molecule::fragmentHasQueryBonds(mol, &filter_e))
      ret |= 2;

   return ret;
}

int MoleculeFingerprintBuilder::_vertex_code (const Graph &graph, int vertex_idx, const void *context)
{
   const Molecule &mol = (const Molecule &)graph;

   if (mol.isPseudoAtom(vertex_idx))
      return CRC32::get(mol.getPseudoAtom(vertex_idx));

   return mol.getAtom(vertex_idx).label;
}

int MoleculeFingerprintBuilder::_edge_code (const Graph &graph, int edge_idx, const void *context)
{
   const Molecule &mol = (const Molecule &)graph;
   return mol.getBond(edge_idx).type;
}

dword MoleculeFingerprintBuilder::_canonicalizeFragment (const Molecule &mol, const Array<int> &vertices, 
         const Array<int> &edges, bool use_atoms, bool use_bonds)
{
   SubgraphHash subgraph_hash(mol);

   if (use_bonds)
      subgraph_hash.cb_edge_code = _edge_code; 
   if (use_atoms)
      subgraph_hash.cb_vertex_code = _vertex_code; 
   subgraph_hash.max_iterations = (edges.size() + 1) / 2;

   return subgraph_hash.getHash(vertices, edges);
}

void MoleculeFingerprintBuilder::_canonicalizeFragmentAndSetBits (const Molecule &mol, const Array<int> &vertices, 
         const Array<int> &edges, bool use_atoms, bool use_bonds, int subgraph_type, dword &bits_set)
{
   bool set_sim = false, set_ord = false, set_any = false, set_tau = false;

   if (subgraph_type == TautomerSuperStructure::ORIGINAL)
   {
      // SIM is made of: rings of size up to 6, trees of size up to 4 edges
      if (use_atoms && use_bonds && !skip_sim && _parameters.sim_qwords > 0)
      {
         set_sim = true;
         if (vertices.size() > 6)
            set_sim = false;
         else if (edges.size() == vertices.size() - 1 && edges.size() > 4)
            set_sim = false;
      }
      
      // ORD and ANY are made of all fragments having more than 2 vertices
      if (vertices.size() > 2)
      {
         if (use_atoms && use_bonds)
         {
            if (!skip_ord && _parameters.ord_qwords > 0)
               set_ord = true;
         }
         else if (_parameters.any_qwords > 0)
         {
            if (use_atoms)
            {
               if (!skip_any_bonds)
                  set_any = true;
            }
            else if (use_bonds)
            {
               if (!skip_any_atoms)
                  set_any = true;
            }
            else if (!skip_any_atoms_bonds)
               set_any = true;
         }
      }
   }

   // TAU is made of fragments without bond types
   if (!use_bonds && !skip_tau && _parameters.tau_qwords > 0)
      set_tau = true;

   if (!set_any && !set_ord && !set_sim && !set_tau)
      return; 

   dword hash = _canonicalizeFragment(mol, vertices, edges, use_atoms, use_bonds);

   if (cb_fragment != 0)
      (*cb_fragment)(mol, vertices, edges, use_atoms, use_bonds, hash);

   // Set bits only if bits_set doesn't have such bits
   dword bits_set_src = bits_set;
   if (!query)
      bits_set_src = 0;

   if (set_sim && !(bits_set_src & 0x01))
   {
      _setBits(hash, getSim(), _parameters.fingerprintSizeSim(), 1);
      bits_set |= 0x01;
   }
   if (set_ord && !(bits_set_src & 0x02))
   {
      _setBits(hash, getOrd(), _parameters.fingerprintSizeOrd(), 2);
      bits_set |= 0x02;
   }
   if (set_any && !(bits_set_src & 0x04))
   {
      _setBits(hash, getAny(), _parameters.fingerprintSizeAny(), 2);
      bits_set |= 0x04;
   }
   if (set_tau && !(bits_set_src & 0x08))
   {
      _setBits(hash, getTau(), _parameters.fingerprintSizeTau(), 2);
      bits_set |= 0x08;
   }
}

void MoleculeFingerprintBuilder::_handleSubgraph (const Graph &graph,
        const Array<int> &vertices, const Array<int> &edges)
{
   const Molecule &mol = (const Molecule &)graph;        

   int subgraph_type;
   if (_tau_super_structure != 0)
      subgraph_type = _tau_super_structure->getSubgraphType(vertices, edges);
   else
      subgraph_type = TautomerSuperStructure::ORIGINAL;
   
   if (subgraph_type == TautomerSuperStructure::NONE)
      return;

   // Check if fragment has query atoms or query bonds
   bool has_query_atoms = Molecule::fragmentHasQueryAtoms(mol, vertices);
   bool has_query_bonds = Molecule::fragmentHasQueryBonds(mol, edges);

   dword bits_set = 0;
   if (!has_query_atoms && !has_query_bonds)
      _canonicalizeFragmentAndSetBits(mol, vertices, edges, true, true, subgraph_type, bits_set);

   dword bits_set_a = bits_set;
   if (!query || (has_query_bonds && !has_query_atoms))
      _canonicalizeFragmentAndSetBits(mol, vertices, edges, true, false, subgraph_type, bits_set_a);

   dword bits_set_b = bits_set;
   if (!query || (!has_query_bonds && has_query_atoms))
      _canonicalizeFragmentAndSetBits(mol, vertices, edges, false, true, subgraph_type, bits_set_b);

   dword bits_set_ab = bits_set_a || bits_set_b;
   if (!query || has_query_atoms || has_query_bonds)
      _canonicalizeFragmentAndSetBits(mol, vertices, edges, false, false, subgraph_type, bits_set_ab);
}

bool MoleculeFingerprintBuilder::_canAtomBeHydrogen (const Molecule &mol, int atom_index)
{
   bool query_atom = false;
   
   if (mol.haveQueryAtoms())   
      if (mol.getQueryAtom(atom_index).type != 0)
         query_atom = true;

   if (!query_atom)
      return mol.getAtom(atom_index).label == ELEM_H;

   const QueryAtom &qa = mol.getQueryAtom(atom_index);

   if (qa.type == QUERY_ATOM_RGROUP)
      // Hydrogen can be represented as R-Group. 
      // Additional check that R-Group can represent hydrogen can be added in 
      // future and in this case RGroups with 2 connectors should be handled carefully.
      return true; 

   if (qa.type == QUERY_ATOM_NOTLIST)
      return !qa.listContains(ELEM_H);
   if (qa.type == QUERY_ATOM_LIST)
      return qa.listContains(ELEM_H);
   
   if (qa.type == QUERY_ATOM_Q || qa.type == QUERY_ATOM_X || qa.type == QUERY_ATOM_A)
      return false;

   throw Error("Unknown query atom type %d", qa.type);
}

void MoleculeFingerprintBuilder::_makeFingerprint (const Molecule &mol)
{
   QS_DEF(Molecule, copy_without_h);
   QS_DEF(Array<int>, vertices);
   int i;

   vertices.clear();

   // remove hydrogens
   for (i = mol.vertexBegin(); i < mol.vertexEnd(); i = mol.vertexNext(i))
      if (!_canAtomBeHydrogen(mol, i))
         vertices.push(i);

   copy_without_h.makeSubmolecule(mol, vertices, 0);

   Obj<TautomerSuperStructure> tau_super_structure;
   Molecule *mol_for_enumeration = &copy_without_h;
   
   if (!query && _parameters.tau_qwords > 0 && !skip_tau)
   {
      tau_super_structure.create();
      tau_super_structure->setBaseMolecule(copy_without_h);
      tau_super_structure->build();

      _tau_super_structure = tau_super_structure.get();
      mol_for_enumeration = &_tau_super_structure->getSuperMolecule(0, 0);
   }
   else
      _tau_super_structure = 0;
   
   if (!skip_ord || !skip_any_atoms || !skip_any_atoms_bonds ||
       !skip_any_bonds || !skip_tau || !skip_sim)
   {
      CycleEnumerator ce(*mol_for_enumeration);
      GraphSubtreeEnumerator se(*mol_for_enumeration);

      bool sim_only = skip_ord && skip_tau && skip_any_atoms &&
                      skip_any_atoms_bonds && skip_any_bonds;

      _is_cycle = true;
      ce.context = this;
      ce.max_length = sim_only ? 6 : 8;
      ce.cb_handle_cycle = _handleCycle;
      ce.process();

      _is_cycle = false;
      se.context = this;
      se.min_vertices = 2;
      se.max_vertices = sim_only ? 5 : 7;
      se.handle_maximal = query;
      se.maximal_critera_value_callback = _maximalSubgraphCriteriaValue;
      se.callback = _handleTree;
      se.process();
   }
   
   if (!skip_ext)
      _calcExtraBits(copy_without_h);
}

void MoleculeFingerprintBuilder::_calcExtraBits (const Molecule &mol)
{
   int counters[9] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
   int i;

   for (i = mol.vertexBegin(); i != mol.vertexEnd(); i = mol.vertexNext(i))
   {
      const Atom &atom = mol.getAtom(i);

      if (!mol.haveQueryAtoms() || mol.getQueryAtom(i).type == 0)
      {
         if (atom.label == ELEM_C)
            counters[0]++;
         else if (atom.label == ELEM_N)
            counters[1]++;
         else if (atom.label == ELEM_O)
            counters[2]++;
         else if (atom.label == ELEM_P)
            counters[3]++;
         else if (atom.label == ELEM_S)
            counters[4]++;
         else if (Element::isHalogen(atom.label))
            counters[5]++;
         else if (atom.label != ELEM_H)
            counters[6]++;
      }
      if (!skip_ext_charge && atom.charge != 0)
         counters[7]++;
      if (atom.isotope != 0)
         counters[8]++;
   }
   
   byte *fp = _total_fingerprint.ptr();

   if (counters[0] > 13) // > 13 C
      fp[0] |= 1;
   if (counters[0] > 16) // > 16 C
      fp[0] |= 2;
   if (counters[0] > 19) // > 19 C
      fp[0] |= 4;
   if (counters[1] > 1) // > 1 N
      fp[0] |= 8;
   if (counters[1] > 2) // > 2 N
      fp[0] |= 16;
   if (counters[2] > 3) // > 3 O
      fp[0] |= 32;
   if (counters[2] > 4) // > 4 O
      fp[0] |= 64;
   if (counters[3] > 0) // have P
      fp[0] |= 128;
   if (counters[4] > 0) // have S
      fp[1] |= 1;
   if (counters[4] > 1) // > 1 S
      fp[1] |= 2;
   if (counters[5] > 1) // > 1 halogen
      fp[1] |= 4;
   if (counters[5] > 2) // > 2 halogen
      fp[1] |= 8;
   if (counters[6] > 0) // have rare atoms
      fp[1] |= 16;
   if (counters[6] > 1) // > 1 rare atom
      fp[1] |= 32;
   if (counters[7] > 0) // have charged atoms
      fp[1] |= 64;
   if (counters[8] > 1) // have isotopes
      fp[1] |= 128;
}

void MoleculeFingerprintBuilder::_setBits (dword hash, byte *fp, int size, int nbits)
{
   unsigned seed = hash;

   // fill random bits
   while (nbits-- > 0)
   {
      seed = seed * 0x8088405 + 1;

      unsigned k = seed % (size * 8);
      unsigned nbyte = k / 8;
      unsigned nbit = k - nbyte * 8;

      fp[nbyte] = fp[nbyte] | (1 << nbit);
   }
}

const byte * MoleculeFingerprintBuilder::get ()
{
   return _total_fingerprint.ptr();
}

byte * MoleculeFingerprintBuilder::getOrd ()
{
   return _total_fingerprint.ptr() + 3;
}

byte * MoleculeFingerprintBuilder::getAny ()
{
   return _total_fingerprint.ptr() + 3 + _parameters.fingerprintSizeOrd();
}

byte * MoleculeFingerprintBuilder::getTau ()
{
   return _total_fingerprint.ptr() + 3 + _parameters.fingerprintSizeOrd() +
          _parameters.fingerprintSizeAny();
}

byte * MoleculeFingerprintBuilder::getSim ()
{
   return _total_fingerprint.ptr() + 3 + _parameters.fingerprintSizeOrd() +
        _parameters.fingerprintSizeAny() + _parameters.fingerprintSizeTau();
}

int MoleculeFingerprintBuilder::countBits_Sim ()
{
   return bitGetOnesCount(getSim(), _parameters.fingerprintSizeSim());
}
